import os
import pickle
import time
import dnnlib
import legacy

import click
import paramiko
import torch
from torch.utils.tensorboard import SummaryWriter

from metrics import metric_main
from tqdm import tqdm


metric = 'fid50k_full'



def eval(
        gan_dir,
        out_dir,
        data_dir,
        t
):
    # make dir =====================================================================
    folder_dir = out_dir
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)
        os.mkdir(os.path.join(folder_dir, 'visualization'))
        print(f'make output dir: {out_dir}')
    
    dataset = 'OASIS' if 'OASIS' in data_dir else 'CANDI'
    # folder_dir = os.path.join(out_dir, f'{dataset}-{t}')

    # if not os.path.exists(folder_dir):
    #     os.mkdir(folder_dir)
    #     # os.mkdir(os.path.join(folder_dir, 'checkpoint'))
    #     os.mkdir(os.path.join(folder_dir, 'visualization'))
    #     # os.mkdir(os.path.join(folder_dir, 'runs'))
    #     print('make dir done!')

    # prevent from renaming a new pkl to best
    with dnnlib.util.open_url(gan_dir) as f:
        snapshot_data = legacy.load_network_pkl(f)
        G_ema = snapshot_data['G_ema'].requires_grad_(False).cuda() # type: ignore
        training_set_kwargs = snapshot_data['training_set_kwargs']
        training_set_kwargs['path'] = data_dir
        # training_set_kwargs['max_size'] = 3694 if eval_on=='data/images' else 2694
        # training_set_kwargs['max_size'] = 6592
        nimg = snapshot_data['nimg']
        del snapshot_data

    G_kwargs = {
        'truncation_psi': t,
    }
    # eval
    result_dict = metric_main.calc_metric(metric=metric, G=G_ema, G_kwargs=G_kwargs,
        dataset_kwargs=training_set_kwargs, num_gpus=1, rank=0, device=torch.device('cuda'))

    fid = result_dict.results.fid50k_full
    total_time = result_dict.total_time_str
    print(f'{t}: {fid}')

    with open(os.path.join(folder_dir, 'test_INFO.txt'), 'a') as log:
        log.write(f'{dataset}-{t}-=================================\n')
        log.write(f'{result_dict}\n')
        log.write(f'{fid:.2f}\n')

if __name__ == '__main__':
    os.environ["CUDA_VISIBLE_DEVICES"] = '2'
    for i in tqdm(range(5,11)):
        eval(
            gan_dir='save/00010-images-mirror-low_shot-kimg25000-batch32-color-translation-cutout/network-snapshot-best.pkl',
            data_dir='data/CANDI-128-256/images',
            out_dir='save/_gan/CANDI',
            t=i/10
        )
    for i in tqdm(range(5,11)):
        eval(
            gan_dir='save/00011-GAN_OASIS-mirror-low_shot-kimg25000-batch32-color-translation-cutout/network-snapshot-best.pkl',
            data_dir='data/OASIS-128-256/images',
            out_dir='save/_gan/OASIS',
            t=i/10
        )